{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c9d89b5c",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/node_postprocessor/ibm_watsonx.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70996d8a",
   "metadata": {},
   "source": [
    "# IBM watsonx.ai\n",
    "\n",
    ">WatsonxRerank is a wrapper for IBM [watsonx.ai](https://www.ibm.com/products/watsonx-ai) Rerank.\n",
    "\n",
    "The aim of these examples is to show how to take advantage of `watsonx.ai` Rerank, Embeddings and LLMs using the `LlamaIndex` postprocessor API."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea35b2b7",
   "metadata": {},
   "source": [
    "## Setting up\n",
    "\n",
    "Install required packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f1fff4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -qU llama-index\n",
    "%pip install -qU llama-index-llms-ibm\n",
    "%pip install -qU llama-index-postprocessor-ibm\n",
    "%pip install -qU llama-index-embeddings-ibm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f406e092",
   "metadata": {},
   "source": [
    "The cell below defines the credentials required to work with watsonx Foundation Models, Embeddings and Rerank.\n",
    "\n",
    "**Action:** Provide the IBM Cloud user API key. For details, see\n",
    "[Managing user API keys](https://cloud.ibm.com/docs/account?topic=account-userapikey&interface=ui)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11d572a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "watsonx_api_key = getpass()\n",
    "os.environ[\"WATSONX_APIKEY\"] = watsonx_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c59782a7",
   "metadata": {},
   "source": [
    "Additionally, you can pass additional secrets as an environment variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f98c573c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"WATSONX_URL\"] = \"your service instance url\"\n",
    "os.environ[\"WATSONX_TOKEN\"] = \"your token for accessing the CPD cluster\"\n",
    "os.environ[\"WATSONX_PASSWORD\"] = \"your password for accessing the CPD cluster\"\n",
    "os.environ[\"WATSONX_USERNAME\"] = \"your username for accessing the CPD cluster\"\n",
    "os.environ[\n",
    "    \"WATSONX_INSTANCE_ID\"\n",
    "] = \"your instance_id for accessing the CPD cluster\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0de9336",
   "metadata": {},
   "source": [
    "**Note**: \n",
    "\n",
    "- To provide context for the API call, you must pass the `project_id` or `space_id`. To get your project or space ID, open your project or space, go to the **Manage** tab, and click **General**. For more information see: [Project documentation](https://www.ibm.com/docs/en/watsonx-as-a-service?topic=projects) or [Deployment space documentation](https://www.ibm.com/docs/en/watsonx/saas?topic=spaces-creating-deployment).\n",
    "- Depending on the region of your provisioned service instance, use one of the urls listed in [watsonx.ai API Authentication](https://ibm.github.io/watsonx-ai-python-sdk/setup_cloud.html#authentication).\n",
    "\n",
    "In this example, we’ll use the `project_id` and Dallas URL."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3ef7659",
   "metadata": {},
   "source": [
    "Provide `PROJECT_ID` that will be used for initialize each watsonx integration instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b933bb7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ID = \"PASTE YOUR PROJECT_ID HERE\"\n",
    "URL = \"https://us-south.ml.cloud.ibm.com\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0308874",
   "metadata": {},
   "source": [
    "## Download data and load documents\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d67b5f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2025-02-24 10:46:16--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8000::154, 2606:50c0:8001::154, 2606:50c0:8002::154, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8000::154|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n",
      "\n",
      "data/paul_graham/pa 100%[===================>]  73,28K  --.-KB/s    in 0,06s   \n",
      "\n",
      "2025-02-24 10:46:17 (1,30 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a65dc946",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd7ac083",
   "metadata": {},
   "source": [
    "## Load the Rerank"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "166b40ca",
   "metadata": {},
   "source": [
    "You might need to adjust rerank parameters for different tasks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15647312",
   "metadata": {},
   "outputs": [],
   "source": [
    "truncate_input_tokens = 512"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a03571f0",
   "metadata": {},
   "source": [
    "#### Initialize `WatsonxRerank` instance."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79dcb01b",
   "metadata": {},
   "source": [
    "You need to specify the `model_id` that will be used for rerank. You can find the list of all the available models in [Supported reranker models](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx#rerank)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33efcbc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor.ibm import WatsonxRerank\n",
    "\n",
    "watsonx_rerank = WatsonxRerank(\n",
    "    model_id=\"cross-encoder/ms-marco-minilm-l-12-v2\",\n",
    "    top_n=2,\n",
    "    url=URL,\n",
    "    project_id=PROJECT_ID,\n",
    "    truncate_input_tokens=truncate_input_tokens,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e79e20d",
   "metadata": {},
   "source": [
    "Alternatively, you can use Cloud Pak for Data credentials. For details, see [watsonx.ai software setup](https://ibm.github.io/watsonx-ai-python-sdk/setup_cpd.html).    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c08ee5a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor.ibm import WatsonxRerank\n",
    "\n",
    "watsonx_rerank = WatsonxRerank(\n",
    "    model_id=\"cross-encoder/ms-marco-minilm-l-12-v2\",\n",
    "    url=URL,\n",
    "    username=\"PASTE YOUR USERNAME HERE\",\n",
    "    password=\"PASTE YOUR PASSWORD HERE\",\n",
    "    instance_id=\"openshift\",\n",
    "    version=\"5.1\",\n",
    "    project_id=PROJECT_ID,\n",
    "    truncate_input_tokens=truncate_input_tokens,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "106479c9",
   "metadata": {},
   "source": [
    "## Load the embedding model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10b86b6b",
   "metadata": {},
   "source": [
    "#### Initialize the `WatsonxEmbeddings` instance.\n",
    "\n",
    ">For more information about `WatsonxEmbeddings` please refer to the sample notebook: <a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/embeddings/ibm_watsonx.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5bf6aaf",
   "metadata": {},
   "source": [
    "You might need to adjust embedding parameters for different tasks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66c3226c",
   "metadata": {},
   "outputs": [],
   "source": [
    "truncate_input_tokens = 512"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf2a9d4f",
   "metadata": {},
   "source": [
    "You need to specify the `model_id` that will be used for embedding. You can find the list of all the available models in [Supported embedding models](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx#embed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c81d51e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.ibm import WatsonxEmbeddings\n",
    "\n",
    "watsonx_embedding = WatsonxEmbeddings(\n",
    "    model_id=\"ibm/slate-30m-english-rtrvr\",\n",
    "    url=URL,\n",
    "    project_id=PROJECT_ID,\n",
    "    truncate_input_tokens=truncate_input_tokens,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8647a00c",
   "metadata": {},
   "source": [
    "Change default settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12645356",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.chunk_size = 512"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1a5570f",
   "metadata": {},
   "source": [
    "#### Build index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb518f53",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents=documents, embed_model=watsonx_embedding\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e36acbef",
   "metadata": {},
   "source": [
    "## Load the LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2924c37",
   "metadata": {},
   "source": [
    "#### Initialize the `WatsonxLLM` instance.\n",
    "\n",
    ">For more information about `WatsonxLLM` please refer to the sample notebook: <a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/ibm_watsonx.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae227410",
   "metadata": {},
   "source": [
    "You need to specify the `model_id` that will be used for inferencing. You can find the list of all the available models in [Supported foundation models](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models.html?context=wx)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba54c1de",
   "metadata": {},
   "source": [
    "You might need to adjust model `parameters` for different models or tasks. For details, refer to [Available MetaNames](https://ibm.github.io/watsonx-ai-python-sdk/fm_model.html#metanames.GenTextParamsMetaNames)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54e38eb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_new_tokens = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "359898de",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.ibm import WatsonxLLM\n",
    "\n",
    "watsonx_llm = WatsonxLLM(\n",
    "    model_id=\"meta-llama/llama-3-3-70b-instruct\",\n",
    "    url=URL,\n",
    "    project_id=PROJECT_ID,\n",
    "    max_new_tokens=max_new_tokens,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55aa40af",
   "metadata": {},
   "source": [
    "## Send a query"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9be7fc6c",
   "metadata": {},
   "source": [
    "#### Retrieve top 10 most relevant nodes, then filter with `WatsonxRerank`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b84f0158",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    llm=watsonx_llm,\n",
    "    similarity_top_k=10,\n",
    "    node_postprocessors=[watsonx_rerank],\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"What did Sam Altman do in this essay?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fdc1eb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final Response: In this essay, Sam Altman was recruited to be the\n",
      "president of Y Combinator (YC), and he agreed to take over the role\n",
      "starting with the winter 2014 batch. He initially declined the offer,\n",
      "wanting to start a startup to make nuclear reactors, but eventually\n",
      "agreed after being persuaded. He began learning the job and taking\n",
      "over responsibilities from the author in the latter part of 2013, and\n",
      "officially took over as president in 2014.\n",
      "______________________________________________________________________\n",
      "Source Node 1/2\n",
      "Node ID: 2ed5d8e7-2681-49b0-a112-ea35cc9a8b9e\n",
      "Similarity: 3.2075154781341553\n",
      "Text: \"You know,\" he said, \"you should make sure Y Combinator isn't\n",
      "the last cool thing you do.\"  At the time I didn't understand what he\n",
      "meant, but gradually it dawned on me that he was saying I should quit.\n",
      "This seemed strange advice, because YC was doing great. But if there\n",
      "was one thing rarer than Rtm offering advice, it was Rtm being wrong.\n",
      "So th...\n",
      "______________________________________________________________________\n",
      "Source Node 2/2\n",
      "Node ID: 6ae17865-aaa7-46a5-bc49-f38abf4a825e\n",
      "Similarity: -1.3127477169036865\n",
      "Text: I asked Jessica if she wanted to be president, but she didn't,\n",
      "so we decided we'd try to recruit Sam Altman. We talked to Robert and\n",
      "Trevor and we agreed to make it a complete changing of the guard. Up\n",
      "till that point YC had been controlled by the original LLC we four had\n",
      "started. But we wanted YC to last for a long time, and to do that it\n",
      "could...\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.response.pprint_utils import pprint_response\n",
    "\n",
    "pprint_response(response, show_source=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baf5df0b",
   "metadata": {},
   "source": [
    "#### Directly retrieve top 2 most similar nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f784254",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    llm=watsonx_llm,\n",
    "    similarity_top_k=2,\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"What did Sam Altman do in this essay?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8465e7a9",
   "metadata": {},
   "source": [
    "Retrieved context is irrelevant and response is hallucinated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dd6af0d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final Response: Sam Altman was one of the founders of the first batch\n",
      "of startups funded by the Summer Founders Program, and he later became\n",
      "the second president of YC.\n",
      "______________________________________________________________________\n",
      "Source Node 1/2\n",
      "Node ID: ba52769a-7342-4e6c-af02-4159216a79a8\n",
      "Similarity: 0.6396056863136902\n",
      "Text: We knew undergrads were deciding then about summer jobs, so in a\n",
      "matter of days we cooked up something we called the Summer Founders\n",
      "Program, and I posted an announcement on my site, inviting undergrads\n",
      "to apply. I had never imagined that writing essays would be a way to\n",
      "get \"deal flow,\" as investors call it, but it turned out to be the\n",
      "perfect ...\n",
      "______________________________________________________________________\n",
      "Source Node 2/2\n",
      "Node ID: 43a6cf9f-8284-45db-bbbd-44109fcb9373\n",
      "Similarity: 0.6334836031239921\n",
      "Text: I wrote this new Lisp, called Bel, in itself in Arc. That may\n",
      "sound like a contradiction, but it's an indication of the sort of\n",
      "trickery I had to engage in to make this work. By means of an\n",
      "egregious collection of hacks I managed to make something close enough\n",
      "to an interpreter written in itself that could actually run. Not fast,\n",
      "but fast enough...\n"
     ]
    }
   ],
   "source": [
    "pprint_response(response, show_source=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
